import numpy as np
import pandas as pd
import talib
from torch.utils.data import Dataset


class XDataset(Dataset):
    def __init__(self, df, index):
        self.df = df
        self.index = index
        self.cols = ["open", "close", "high", "low"]
        self.extened = False
        self.batch_normal_cols = None
        self.history = 40
        self.init()

    def init(self):
        df = self.df
        cols = self.cols

        if "tp" in df.columns:
            return

        df["clv"] = ((df["close"] - df["low"]) - (df["high"] - df["close"])) / (
            df["high"] - df["low"] + 1e-9
        )
        cols.append("clv")

        df["tp"] = (df["high"] + df["low"] + df["close"]) / 3
        cols.append("tp")

        for i in [7, 12, 25]:
            name = f"ma_d{i}"
            ma = talib.MA(df["tp"], timeperiod=i)
            ma = pd.Series(ma, index=df.index)
            ma_d = ma - ma.shift(1)
            df[name] = ma_d
            cols.append(name)

        df["rsi"] = talib.RSI(df["tp"], timeperiod=24) / 100
        cols.append("rsi")
        df["adx"] = talib.ADX(df["high"], df["low"], df["close"], timeperiod=24) / 100
        cols.append("adx")

    def __getitem__(self, item):
        if self.extened:
            return self.df.loc[self.index[item], self.batch_normal_cols].to_numpy()
        history = self.history
        normal_cols = ["open", "high", "low", "close", "tp"]
        seqcols = list(set(self.cols[: self.cols.index("rsi")]).difference(normal_cols))
        othercols = self.cols[self.cols.index("rsi") :]
        results = []
        origin_idx = self.df.index.get_loc(self.index[item])

        seqmain = (
            self.df[normal_cols]
            .iloc[origin_idx - history + 1 : origin_idx + 1]
            .to_numpy()
        )
        seqmain = (seqmain - seqmain.min()) / (seqmain.max() - seqmain.min())
        if seqcols:
            seqadd = (
                self.df[seqcols]
                .iloc[origin_idx - history + 1 : origin_idx + 1]
                .to_numpy()
            )
            seqmain = np.concatenate([seqmain, seqadd], axis=1)
        results.append(seqmain.flatten())
        if othercols:
            results.append(self.df[othercols].iloc[origin_idx].to_numpy().flatten())
        out = np.concatenate(results, axis=0)
        return out

    def __array__(self):
        if self.extened:
            return self.df.loc[self.index, self.batch_normal_cols].to_numpy()
        history = self.history
        normal_cols = ["open", "high", "low", "close", "tp"]
        seqcols = list(set(self.cols[: self.cols.index("rsi")]).difference(normal_cols))
        othercols = self.cols[self.cols.index("rsi") :]
        batch_normal_cols = []

        need_concats = []
        for i in range(history - 1, -1, -1):
            for name in normal_cols + seqcols:
                fname = f"{name}_x{i}"
                need_concats.append(self.df[name].shift(i))
                batch_normal_cols.append(fname)
        old_columns = self.df.columns
        self.df = pd.concat([self.df, *need_concats], axis=1)
        self.df.columns = [*old_columns, *batch_normal_cols]
        consider_value_norms = []
        for name in batch_normal_cols:
            for prefix in normal_cols:
                if name.startswith(prefix):
                    consider_value_norms.append(name)

        min_price = self.df[consider_value_norms].min(axis=1).to_numpy()
        max_price = self.df[consider_value_norms].max(axis=1).to_numpy()
        self.df[consider_value_norms] = (
            self.df[consider_value_norms].to_numpy() - min_price[:, np.newaxis]
        ) / (max_price[:, np.newaxis] - min_price[:, np.newaxis] + 1e-9)

        batch_normal_cols.extend(othercols)
        self.extened = True
        self.batch_normal_cols = batch_normal_cols
        return self.df.loc[self.index, self.batch_normal_cols].to_numpy()

    def __len__(self):
        return len(self.index)

    def __iter__(self):
        for i in range(len(self)):
            yield self[i]


class YDataset(Dataset):
    def __init__(self, df, index):
        self.df = df
        self.index = index
        self.init()

    def init(self):
        df = self.df
        delta = 1
        if "tag" in df.columns:
            return
        df["tag"] = (df["close"].shift(-delta) - df["close"]) / df["close"]
        # df["tag"] = np.where(
        #     df["tag"] > 0.001,
        #     2,
        #     np.where(
        #         df["tag"] < -0.001,
        #         0, 1
        #     )
        # )

        # smoothed = gaussian_filter1d(df["close"], sigma=1.4)
        # smoothed = pd.Series(smoothed, index=df.index)
        # diff1 = smoothed - smoothed.shift(1)
        # diffpos = (diff1.shift(-1) >= 0) & (diff1.shift(0) < 0)
        # diffneg = (diff1.shift(-1) <= 0) & (diff1.shift(0) > 0)
        #
        # diff_choose = df[diffpos | diffneg]
        # diff_change = (diff_choose["close"].shift(-1) - diff_choose["close"]) / (diff_choose["close"] + 1e-9)
        # diff_change = diff_change[np.abs(diff_change) > 0.005]
        #
        # diffpos[df.index.difference(diff_change.index)] = False
        # diffneg[df.index.difference(diff_change.index)] = False
        # df["peak_change"] = diff_change
        # # df["tag"] = 0.0
        # #
        # # next_value = None
        # # for idx in df.index[::-1]:
        # #     if np.isnan(df.loc[idx, "peak_change"]):
        # #         if next_value is None:
        # #             continue
        # #         df.loc[idx, "tag"] = (next_value - df.loc[idx, "close"]) / df.loc[idx, "close"]
        # #     else:
        # #         df.loc[idx, "tag"] = df.loc[idx, "peak_change"]
        # #         next_value = df.loc[idx, "close"]
        #
        #
        # df["peak_change"] = np.where(
        #     np.isnan(df["peak_change"]),
        #     0,
        #     df["peak_change"]
        # )
        # df["tag"] = np.where(
        #     df["peak_change"] > 0.005,
        #     2,
        #     np.where(
        #         df["peak_change"] < -0.005,
        #         0, 1
        #     )
        # )
        #
        # # df["tag"] = df["tag"] * 100
        return df["tag"].to_numpy()

    def __getitem__(self, item):
        return self.df.loc[self.index[item], "tag"]

    def __len__(self):
        return len(self.index)

    def __array__(self):
        return self.df.loc[self.index, "tag"].to_numpy()

    def __iter__(self):
        for i in range(len(self)):
            yield self[i]


class TwoDataset(Dataset):
    def __init__(self, X, y):
        self.X = X
        self.y = y

    def __getitem__(self, item):
        return self.X[item], self.y[item]

    def __len__(self):
        return len(self.X)
